import java.util.*;

public class T0311_HuffmanCoding_Demo {

    // 定义霍夫曼树的节点类
    static class Node {
        String word;
        int frequency;
        Node left;
        Node right;

        Node(String word, int frequency) {
            this.word = word;
            this.frequency = frequency;
        }
    }

    // 构建霍夫曼树
    public static Node buildHuffmanTree(Map<String, Integer> wordFrequencies) {
        PriorityQueue<Node> priorityQueue = new PriorityQueue<>(Comparator.comparingInt(node -> node.frequency));

        // 将每个词语及其频率作为一个节点，放入优先队列
        for (Map.Entry<String, Integer> entry : wordFrequencies.entrySet()) {
            priorityQueue.add(new Node(entry.getKey(), entry.getValue()));
        }

        // 构建霍夫曼树
        while (priorityQueue.size() > 1) {
            Node left = priorityQueue.poll();
            Node right = priorityQueue.poll();
            Node parent = new Node(null, left.frequency + right.frequency);
            parent.left = left;
            parent.right = right;
            priorityQueue.add(parent);
        }

        // 返回霍夫曼树的根节点
        return priorityQueue.poll();
    }

    // 生成编码表
    public static void generateCodes(Node node, String code, Map<String, String> huffmanCodes) {
        if (node == null) {
            return;
        }
        if (node.word != null) {
            huffmanCodes.put(node.word, code);
        }
        generateCodes(node.left, code + "0", huffmanCodes);
        generateCodes(node.right, code + "1", huffmanCodes);
    }

    // 霍夫曼编码主方法
    public static Map<String, String> huffmanCoding(Map<String, Integer> wordFrequencies) {
        Node root = buildHuffmanTree(wordFrequencies);
        Map<String, String> huffmanCodes = new HashMap<>();
        generateCodes(root, "", huffmanCodes);
        return huffmanCodes;
    }

    // 测试方法
    public static void main(String[] args) {
        Map<String, Integer> wordFrequencies = new HashMap<>();
        wordFrequencies.put("a", 5);
        wordFrequencies.put("b", 9);
        wordFrequencies.put("c", 12);
        wordFrequencies.put("d", 13);
        wordFrequencies.put("e", 16);
        wordFrequencies.put("f", 45);

        Map<String, String> huffmanCodes = huffmanCoding(wordFrequencies);
        System.out.println("霍夫曼编码结果：");
        for (Map.Entry<String, String> entry : huffmanCodes.entrySet()) {
            System.out.println(entry.getKey() + ": " + entry.getValue());
        }
    }
}
